package inference;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

import model.ModelParameters;
import model.PredicateInfo;
import model.TableInformation;
import utils.IntPair;
import utils.Pair;

import com.analog.lyric.dimple.model.Discrete;

public class RootFactor {
	public static Pair<int[][], double[]> createRootTopicFactor(Discrete table_topic_type_variable, Discrete table_topic_variable,
			   ModelParameters params) {
		Object[] topic_types = table_topic_type_variable.getDomain().getElements();
		int unk_topic_type_id = topic_types.length - 1;
		Object[] topics = table_topic_variable.getDomain().getElements();
		int unk_topic_id = topics.length - 1;
		
		HashSet<Integer> type_set = new HashSet<Integer>();
		for (Object type : topic_types) {
			type_set.add((Integer)type);
		}
		
		// Set up output data structures.
		int init_size = topic_types.length * topics.length;
		ArrayList<IntPair> indices = new ArrayList<IntPair>(init_size);
		ArrayList<Double> probabilities = new ArrayList<Double>(init_size);
		
		TableInformation table_info = params.getTableInformation();

		HashMap<Integer, HashSet<Integer> > type_to_topics = new HashMap<Integer, HashSet<Integer> >();
		for (int t = 0; t < topics.length; t++) {
			if (params.useUnknownTopic() && t == unk_topic_id) {
				continue;
			}
			Integer topic_id = (Integer)(topics[t]);
			HashSet<Integer> types = table_info.getTypesForEntityId(topic_id);
			System.out.println(table_info.getStringForId(topic_id));
			if (types != null) {
				for (Integer type : types) {
					if (type_set.contains(type)) {
						HashSet<Integer> type_topics = type_to_topics.get(type);
						if (type_topics == null) {
							type_topics = new HashSet<Integer>();
							type_to_topics.put(type, type_topics);
						}
						type_topics.add(t);
					}
				}	
			}
		}
		
		for (int t = 0; t < topic_types.length; t++) {
			if (t == unk_topic_type_id) {
				for (int i = 0; i < topics.length; i++) {
					indices.add(new IntPair(t, i));
					probabilities.add(1.0 / topics.length);
				}
			} else {
				Integer topic_type_id = (Integer)(topic_types[t]);
				HashSet<Integer> type_topics = type_to_topics.get(topic_type_id);
				if (type_topics == null) {
					if (params.useUnknownTopic()) {
						indices.add(new IntPair(t, unk_topic_id));
						probabilities.add(1.0);
					} else {
						double prob = 1.0 / topics.length;
						for (int i = 0; i < topics.length; i++) {
							indices.add(new IntPair(t, i));
							probabilities.add(prob);
						}
					}
				} else {
					if (params.useUnknownTopic()) {
						for (Integer topic_id : type_topics) {
							indices.add(new IntPair(t, topic_id));
							probabilities.add(1.0 / (type_topics.size() + 1));
						}
						indices.add(new IntPair(t, unk_topic_id));
						probabilities.add(1.0 / (type_topics.size() + 1));
					} else {
						double prob = 1.0 / type_topics.size();
						for (Integer topic_id : type_topics) {
							indices.add(new IntPair(t, topic_id));
							probabilities.add(prob);
						}
					}
				}
			}
		}
		
		return translateListsToPrimitives(indices, probabilities);
	}
	
	public static Pair<int[][], double[]> createRootTopicFactorReverse(Discrete table_topic_type_variable, Discrete table_topic_variable,
			   ModelParameters params) {
		Object[] topic_types = table_topic_type_variable.getDomain().getElements();
		int unk_topic_type_id = topic_types.length - 1;
		Object[] topics = table_topic_variable.getDomain().getElements();
		int unk_topic_id = topics.length - 1;
		
		HashSet<Integer> type_set = new HashSet<Integer>();
		for (Object type : topic_types) {
			type_set.add((Integer)type);
		}
		
		// Set up output data structures.
		int init_size = topic_types.length * topics.length;
		ArrayList<IntPair> indices = new ArrayList<IntPair>(init_size);
		ArrayList<Double> probabilities = new ArrayList<Double>(init_size);
		
		TableInformation table_info = params.getTableInformation();
		for (int t = 0; t < topics.length; t++) {
			HashSet<Integer> types = table_info.getTypesForEntityId((Integer)(topics[t]));
			if (types == null) {
				indices.add(new IntPair(t, unk_topic_type_id));
				probabilities.add(1.0);
			} else {
				for (int i = 0; i < topic_types.length; i++) {
					Integer type = (Integer)(topic_types[i]);
					if (types.contains(type)) {
						indices.add(new IntPair(t, i));
						probabilities.add(1.0 / types.size());
					}
				}
			}
		}
		
		return translateListsToPrimitives(indices, probabilities);
	}
	
	public static Pair<int[][], double[]> createRootPredicateFactor(Discrete table_topic_type_variable, Discrete predicate_variable,
			   ModelParameters params) {
		Object[] predicates = predicate_variable.getDomain().getElements();
		int unk_pred_idx = predicates.length - 1;
		Object[] topic_types = table_topic_type_variable.getDomain().getElements();
		int unk_topic_type_idx = topic_types.length - 1;
		
		// Set up output data structures.
		int init_size = predicates.length;
		ArrayList<IntPair> indices = new ArrayList<IntPair>(init_size);
		ArrayList<Double> probabilities = new ArrayList<Double>(init_size);
		
		TableInformation table_info = params.getTableInformation();

		HashMap<Integer, HashSet<Integer> > domain_to_preds = new HashMap<Integer, HashSet<Integer> >();
		for (int p = 0; p < predicates.length; p++) {
			if (p == unk_pred_idx) {
				continue;
			}
			Integer predicate_id = (Integer)(predicates[p]);
			PredicateInfo pred_info = table_info.getPredicateInfo(predicate_id);
			Integer domain_id = pred_info.getDomainId();
			HashSet<Integer> preds = domain_to_preds.get(domain_id);
			if (preds == null) {
				preds = new HashSet<Integer>();
				domain_to_preds.put(domain_id, preds);
			}
			preds.add(p);
		}
		
		for (int t = 0; t < topic_types.length; t++) {
			if (t == unk_topic_type_idx) {
				indices.add(new IntPair(unk_topic_type_idx, unk_pred_idx));
				probabilities.add(1.0);
				continue;
			}
			Integer topic_type_id = (Integer)(topic_types[t]);
			HashSet<Integer> preds = domain_to_preds.get(topic_type_id);
			if (preds == null) {
				indices.add(new IntPair(t, unk_pred_idx));
				probabilities.add(1.0);
			} else {
				for (Integer pred_idx : preds) {
					indices.add(new IntPair(t, pred_idx));
					probabilities.add(1.0 / (preds.size() + 1));
				}
				indices.add(new IntPair(t, unk_pred_idx));
				probabilities.add(1.0 / (preds.size() + 1));
			}
		}
		
		return translateListsToPrimitives(indices, probabilities);
	}
	
	private static Pair<int[][], double[]> translateListsToPrimitives(ArrayList<IntPair> indices, ArrayList<Double> probabilities) {
		int[][] indices_arr = new int[indices.size()][2];
		double[] prob_arr = new double[probabilities.size()];
		
		Iterator<IntPair> it = indices.iterator();
		Iterator<Double> it_prob = probabilities.iterator();
		for (int i = 0; i < indices.size(); i++) {
			IntPair next = it.next();
			indices_arr[i] = next.toArray();
			prob_arr[i] = it_prob.next().doubleValue();
//			System.out.println(next.toString() + " " + prob_arr[i]);
		}
		
		return new Pair<int[][], double[]>(indices_arr, prob_arr);
	}
	
	private static int indexOfEntity(Object[] entities, Integer entity) {
		for (int i = 0; i < entities.length; i++) {
			if (((Integer)(entities[i])).equals(entity)) return i;
		}
		return -1;
	}
}
